{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This mounts your Google Drive to the Colab VM.\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "\n",
    "# TODO: Enter the foldername in your Drive where you have saved the unzipped\n",
    "# assignment folder, e.g. 'cs231n/assignments/assignment3/'\n",
    "FOLDERNAME = None\n",
    "assert FOLDERNAME is not None, \"[!] Enter the foldername.\"\n",
    "\n",
    "# Now that we've mounted your Drive, this ensures that\n",
    "# the Python interpreter of the Colab VM can load\n",
    "# python files from within it.\n",
    "import sys\n",
    "sys.path.append('/content/drive/My Drive/{}'.format(FOLDERNAME))\n",
    "\n",
    "# This downloads the COCO dataset to your Drive\n",
    "# if it doesn't already exist.\n",
    "%cd /content/drive/My\\ Drive/$FOLDERNAME/cs231n/datasets/\n",
    "!bash get_datasets.sh\n",
    "%cd /content/drive/My\\ Drive/$FOLDERNAME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using GPU\n",
    "\n",
    "Go to `Runtime > Change runtime type` and set `Hardware accelerator` to `GPU`. This will reset Colab. **Rerun the top cell to mount your Drive again.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rXtwio7U3mSI",
    "tags": [
     "pdf-title"
    ]
   },
   "source": [
    "# Self-Supervised Learning \n",
    "\n",
    "## What is self-supervised learning?\n",
    "Modern day machine learning requires lots of labeled data. But often times it's challenging and/or expensive to obtain large amounts of human-labeled data. Is there a way we could ask machines to automatically learn a model which can generate good visual representations without a labeled dataset? Yes, enter self-supervised learning! \n",
    "\n",
    "Self-supervised learning (SSL) allows models to automatically learn a \"good\" representation space using the data in a given dataset without the need for their labels. Specifically, if our dataset were a bunch of images, then self-supervised learning allows a model to learn and generate a \"good\" representation vector for images. \n",
    "\n",
    "The reason SSL methods have seen a surge in popularity is because the learnt model continues to perform well on other datasets as well i.e. new datasets on which the model was not trained on!\n",
    "\n",
    "## What makes a \"good\" representation?\n",
    "A \"good\" representation vector needs to capture the important features of the image as it relates to the rest of the dataset. This means that images in the dataset representing semantically similar entities should have similar representation vectors, and different  images in the dataset should have different representation vectors. For example, two images of an apple should have similar representation vectors, while an image of an apple and an image of a banana should have different representation vectors.\n",
    "\n",
    "## Contrastive Learning: SimCLR\n",
    "Recently, [SimCLR](https://arxiv.org/pdf/2002.05709.pdf) introduces a new architecture which uses **contrastive learning** to learn good visual representations. Contrastive learning aims to learn similar representations for similar images and different representations for different images. As we will see in this notebook, this simple idea allows us to train a surprisingly good model without using any labels.\n",
    "\n",
    "Specifically, for each image in the dataset, SimCLR generates two differently augmented views of that image, called a **positive pair**. Then, the model is encouraged to generate similar representation vectors for this pair of images. See below for an illustration of the architecture (Figure 2 from the paper)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this cell to view the SimCLR architecture.\n",
    "from IPython.display import Image\n",
    "Image('images/simclr_fig2.png', width=500)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given an image **x**, SimCLR uses two different data augmentation schemes **t** and **t'** to generate the positive pair of images **$\\hat{x}_i$** and **$\\hat{x}_j$**. $f$ is a basic encoder net that  extracts representation vectors from the augmented data samples, which yields **$h_i$** and **$h_j$**, respectively. Finally, a small neural network projection head $g$ maps the representation vectors to the space where the contrastive loss is applied. The goal of the contrastive loss is to maximize agreement between the final vectors **$z_i = g(h_i)$** and **$z_j = g(h_j)$**. We will discuss the contrastive loss in more detail later, and you will get to implement it.\n",
    "\n",
    "After training is completed, we throw away the projection head $g$ and only use $f$ and the representation $h$ to perform downstream tasks, such as classification. You will get a chance to finetune a layer on top of a trained SimCLR model for a classification task and compare its performance with a baseline model (without self-supervised learning)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pretrained Weights\n",
    "For your convenience, we have given you pretrained weights (trained for ~18 hours on CIFAR-10) for the SimCLR model. Run the following cell to download pretrained model weights to be used later. (This will take ~1 minute)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "DIR=pretrained_model/\n",
    "if [ ! -d \"$DIR\" ]; then\n",
    "    mkdir \"$DIR\" \n",
    "fi\n",
    "\n",
    "URL=http://downloads.cs.stanford.edu/downloads/cs231n/pretrained_simclr_model.pth\n",
    "FILE=pretrained_model/pretrained_simclr_model.pth\n",
    "if [ ! -f \"$FILE\" ]; then\n",
    "    echo \"Downloading weights...\"\n",
    "    wget \"$URL\" -O \"$FILE\"\n",
    "fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup cell.\n",
    "%pip install thop\n",
    "import torch\n",
    "import os\n",
    "import importlib\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch.optim as optim\n",
    "import torch.nn as nn\n",
    "import random\n",
    "from thop import profile, clever_format\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.datasets import CIFAR10\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7Kssxyfx3mSP"
   },
   "source": [
    "# Data Augmentation\n",
    "\n",
    "Our first step is to perform data augmentation. Implement the `compute_train_transform()` function in `cs231n/simclr/data_utils.py` to apply the following random transformations:\n",
    "\n",
    "1. Randomly resize and crop to 32x32.\n",
    "2. Horizontally flip the image with probability 0.5\n",
    "3. With a probability of 0.8, apply color jitter (see `compute_train_transform()` for definition)\n",
    "4. With a probability of 0.2, convert the image to grayscale"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now complete `compute_train_transform()` and `CIFAR10Pair.__getitem__()` in `cs231n/simclr/data_utils.py` to apply the data augmentation transform and generate **$\\hat{x}_i$** and **$\\hat{x}_j$**.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Fx_VMEeT3mSP"
   },
   "source": [
    "Test to make sure that your data augmentation code is correct:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 282,
     "status": "ok",
     "timestamp": 1618984254919,
     "user": {
      "displayName": "Sean Jeng Liu",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjAZgcQiLo3RKO3ZgIxMOZGW1pIJtlPubTTjpaA=s64",
      "userId": "12964162795476174470"
     },
     "user_tz": 240
    },
    "id": "9cBFOd_A3mSQ"
   },
   "outputs": [],
   "source": [
    "from cs231n.simclr.data_utils import *\n",
    "from cs231n.simclr.contrastive_loss import *\n",
    "\n",
    "answers = torch.load('simclr_sanity_check.key')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import torchvision\n",
    "from torchvision.datasets import CIFAR10\n",
    "\n",
    "def test_data_augmentation(correct_output=None):\n",
    "    train_transform = compute_train_transform(seed=2147483647)\n",
    "    trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=train_transform)\n",
    "    trainloader = torch.utils.data.DataLoader(trainset, batch_size=2, shuffle=False, num_workers=2)\n",
    "    dataiter = iter(trainloader)\n",
    "    images, labels = dataiter.next()\n",
    "    img = torchvision.utils.make_grid(images)\n",
    "    img = img / 2 + 0.5     # unnormalize\n",
    "    npimg = img.numpy()\n",
    "    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n",
    "    plt.show()\n",
    "    output = images\n",
    "    \n",
    "    print(\"Maximum error in data augmentation: %g\"%rel_error( output.numpy(), correct_output.numpy()))\n",
    "\n",
    "# Should be less than 1e-07.\n",
    "test_data_augmentation(answers['data_augmentation'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i76TPfFL3mSQ"
   },
   "source": [
    "# Base Encoder and Projection Head\n",
    "The next steps are to apply the base encoder and projection head to the augmented samples **$\\hat{x}_i$** and **$\\hat{x}_j$**.\n",
    "\n",
    "The base encoder $f$ extracts representation vectors for the augmented samples. The SimCLR paper found that using deeper and wider models improved performance and thus chose [ResNet](https://arxiv.org/pdf/1512.03385.pdf) to use as the base encoder. The output of the base encoder are the representation vectors **$h_i = f(\\hat{x}_i$)** and **$h_j = f(\\hat{x}_j$)**.\n",
    "\n",
    "The projection head $g$ is a small neural network that maps the representation vectors **$h_i$** and **$h_j$** to the space where the contrastive loss is applied. The paper found that using a nonlinear projection head improved the representation quality of the layer before it. Specifically, they used a MLP with one hidden layer as the projection head $g$. The contrastive loss is then computed based on the outputs **$z_i = g(h_i$)** and **$z_j = g(h_j$)**.\n",
    "\n",
    "We provide implementations of these two parts in `cs231n/simclr/model.py`. Please skim through the file and make sure you understand the implementation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z_9ORIPN3mSQ"
   },
   "source": [
    "# SimCLR: Contrastive Loss\n",
    "\n",
    "A mini-batch of $N$ training images yields a total of $2N$ data-augmented examples. For each positive pair $(i, j)$ of augmented examples, the contrastive loss function aims to maximize the agreement of vectors $z_i$ and $z_j$. Specifically, the loss is the normalized temperature-scaled cross entropy loss and aims to maximize the agreement of $z_i$ and $z_j$ relative to all other augmented examples in the batch:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "l \\; (i, j) = -\\log \\frac{\\exp (\\;\\text{sim}(z_i, z_j)\\; / \\;\\tau) }{\\sum_{k=1}^{2N} \\mathbb{1}_{k \\neq i} \\exp (\\;\\text{sim} (z_i, z_k) \\;/ \\;\\tau) }\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where $\\mathbb{1} \\in \\{0, 1\\}$ is an indicator function that outputs $1$ if $k\\neq i$ and $0$ otherwise. $\\tau$ is a temperature parameter that determines how fast the exponentials increase.\n",
    "\n",
    "sim$(z_i, z_j) = \\frac{z_i \\cdot z_j}{|| z_i || || z_j ||}$ is the (normalized) dot product between vectors $z_i$ and $z_j$. The higher the similarity between $z_i$ and $z_j$, the larger the dot product is, and the larger the numerator becomes. The denominator normalizes the value by summing across $z_i$ and all other augmented examples $k$ in the batch. The range of the normalized value is $(0, 1)$, where a high score close to $1$ corresponds to a high similarity between the positive pair $(i, j)$ and low similarity between $i$ and other augmented examples $k$ in the batch. The negative log then maps the range $(0, 1)$ to the loss values $(\\inf, 0)$. \n",
    "\n",
    "The total loss is computed across all positive pairs $(i, j)$ in the batch. Let $z = [z_1, z_2, ..., z_{2N}]$ include all the augmented examples in the batch, where $z_{1}...z_{N}$ are outputs of the left branch, and $z_{N+1}...z_{2N}$ are outputs of the right branch. Thus, the positive pairs are $(z_{k}, z_{k + N})$ for $\\forall k \\in [1, N]$. \n",
    "\n",
    "Then, the total loss $L$ is:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$\n",
    "L = \\frac{1}{2N} \\sum_{k=1}^N [ \\; l(k, \\;k+N) + l(k+N, \\;k)\\;]\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**NOTE:** this equation is slightly different from the one in the paper. We've rearranged the ordering of the positive pairs in the batch, so the indices are different. The rearrangement makes it easier to implement the code in vectorized form.\n",
    "\n",
    "We'll walk through the steps of implementing the loss function in vectorized form. Implement the functions `sim`, `simclr_loss_naive` in `cs231n/simclr/contrastive_loss.py`. Test your code by running the sanity checks below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "executionInfo": {
     "elapsed": 269,
     "status": "ok",
     "timestamp": 1618984257779,
     "user": {
      "displayName": "Sean Jeng Liu",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjAZgcQiLo3RKO3ZgIxMOZGW1pIJtlPubTTjpaA=s64",
      "userId": "12964162795476174470"
     },
     "user_tz": 240
    },
    "id": "b-CiVKoe3mSR"
   },
   "outputs": [],
   "source": [
    "from cs231n.simclr.contrastive_loss import *\n",
    "answers = torch.load('simclr_sanity_check.key')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 377,
     "status": "ok",
     "timestamp": 1618984258715,
     "user": {
      "displayName": "Sean Jeng Liu",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjAZgcQiLo3RKO3ZgIxMOZGW1pIJtlPubTTjpaA=s64",
      "userId": "12964162795476174470"
     },
     "user_tz": 240
    },
    "id": "9VMpPjXc3mSS",
    "outputId": "a2f68080-33d7-4994-bb7e-3d89ba52b4aa"
   },
   "outputs": [],
   "source": [
    "def test_sim(left_vec, right_vec, correct_output):\n",
    "    output = sim(left_vec, right_vec).cpu().numpy()\n",
    "    print(\"Maximum error in sim: %g\"%rel_error(correct_output.numpy(), output))\n",
    "\n",
    "# Should be less than 1e-07.\n",
    "test_sim(answers['left'][0], answers['right'][0], answers['sim'][0])\n",
    "test_sim(answers['left'][1], answers['right'][1], answers['sim'][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 345,
     "status": "ok",
     "timestamp": 1618984260903,
     "user": {
      "displayName": "Sean Jeng Liu",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjAZgcQiLo3RKO3ZgIxMOZGW1pIJtlPubTTjpaA=s64",
      "userId": "12964162795476174470"
     },
     "user_tz": 240
    },
    "id": "SMpreOIw3mSS",
    "outputId": "94ef3282-98f8-43ee-d5ca-b303d477a10a"
   },
   "outputs": [],
   "source": [
    "def test_loss_naive(left, right, tau, correct_output):\n",
    "    naive_loss = simclr_loss_naive(left, right, tau).item()\n",
    "    print(\"Maximum error in simclr_loss_naive: %g\"%rel_error(correct_output, naive_loss))\n",
    "\n",
    "# Should be less than 1e-07.\n",
    "test_loss_naive(answers['left'], answers['right'], 5.0, answers['loss']['5.0'])\n",
    "test_loss_naive(answers['left'], answers['right'], 1.0, answers['loss']['1.0'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_cEoU6hM3mSS"
   },
   "source": [
    "Now implement the vectorized version by implementing `sim_positive_pairs`, `compute_sim_matrix`, `simclr_loss_vectorized` in `cs231n/simclr/contrastive_loss.py`.  Test your code by running the sanity checks below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 522,
     "status": "ok",
     "timestamp": 1618984262706,
     "user": {
      "displayName": "Sean Jeng Liu",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjAZgcQiLo3RKO3ZgIxMOZGW1pIJtlPubTTjpaA=s64",
      "userId": "12964162795476174470"
     },
     "user_tz": 240
    },
    "id": "uW4g-aB_3mSS",
    "outputId": "36676e11-f75d-49c6-ec8a-ca8a321eff9f"
   },
   "outputs": [],
   "source": [
    "def test_sim_positive_pairs(left, right, correct_output):\n",
    "    sim_pair = sim_positive_pairs(left, right).cpu().numpy()\n",
    "    print(\"Maximum error in sim_positive_pairs: %g\"%rel_error(correct_output.numpy(), sim_pair))\n",
    "\n",
    "# Should be less than 1e-07.\n",
    "test_sim_positive_pairs(answers['left'], answers['right'], answers['sim'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 310,
     "status": "ok",
     "timestamp": 1618984263449,
     "user": {
      "displayName": "Sean Jeng Liu",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjAZgcQiLo3RKO3ZgIxMOZGW1pIJtlPubTTjpaA=s64",
      "userId": "12964162795476174470"
     },
     "user_tz": 240
    },
    "id": "Fm5z3N7k3mST",
    "outputId": "a91d80b9-02c1-4728-b11d-1302897efbf3"
   },
   "outputs": [],
   "source": [
    "def test_sim_matrix(left, right, correct_output):\n",
    "    out = torch.cat([left, right], dim=0)\n",
    "    sim_matrix = compute_sim_matrix(out).cpu()\n",
    "    assert torch.isclose(sim_matrix, correct_output).all(), \"correct: {}. got: {}\".format(correct_output, sim_matrix)\n",
    "    print(\"Test passed!\")\n",
    "\n",
    "test_sim_matrix(answers['left'], answers['right'], answers['sim_matrix'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 290,
     "status": "ok",
     "timestamp": 1618984265273,
     "user": {
      "displayName": "Sean Jeng Liu",
      "photoUrl": "https://lh3.googleusercontent.com/a-/AOh14GjAZgcQiLo3RKO3ZgIxMOZGW1pIJtlPubTTjpaA=s64",
      "userId": "12964162795476174470"
     },
     "user_tz": 240
    },
    "id": "f5Eps9Pc3mST",
    "outputId": "ad26a099-9d1a-4021-96f1-9e9e3c3e1ad4"
   },
   "outputs": [],
   "source": [
    "def test_loss_vectorized(left, right, tau, correct_output):\n",
    "    vec_loss = simclr_loss_vectorized(left, right, tau, device).item()\n",
    "    print(\"Maximum error in loss_vectorized: %g\"%rel_error(correct_output, vec_loss))\n",
    "\n",
    "# Should be less than 1e-07.\n",
    "test_loss_vectorized(answers['left'], answers['right'], 5.0, answers['loss']['5.0'])\n",
    "test_loss_vectorized(answers['left'], answers['right'], 1.0, answers['loss']['1.0'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implement the train function\n",
    "Complete the `train()` function in `cs231n/simclr/utils.py` to obtain the model's output and use `simclr_loss_vectorized` to compute the loss. (Please take a look at the `Model` class in `cs231n/simclr/model.py` to understand the model pipeline and the returned values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cs231n.simclr.data_utils import *\n",
    "from cs231n.simclr.model import *\n",
    "from cs231n.simclr.utils import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train the SimCLR model\n",
    "\n",
    "Run the following cells to load in the pretrained weights and continue to train a little bit more. This part will take ~10 minutes and will output to `pretrained_model/trained_simclr_model.pth`.\n",
    "\n",
    "**NOTE:** Don't worry about logs such as '_[WARN] Cannot find rule for ..._'. These are related to another module used in the notebook. You can verify the integrity of your code changes through our provided prompts and comments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uHXYbTba3mSU"
   },
   "outputs": [],
   "source": [
    "# Do not modify this cell.\n",
    "feature_dim = 128\n",
    "temperature = 0.5\n",
    "k = 200\n",
    "batch_size = 64\n",
    "epochs = 1\n",
    "temperature = 0.5\n",
    "percentage = 0.5\n",
    "pretrained_path = './pretrained_model/pretrained_simclr_model.pth'\n",
    "\n",
    "# Prepare the data.\n",
    "train_transform = compute_train_transform()\n",
    "train_data = CIFAR10Pair(root='data', train=True, transform=train_transform, download=True)\n",
    "train_data = torch.utils.data.Subset(train_data, list(np.arange(int(len(train_data)*percentage))))\n",
    "train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=16, pin_memory=True, drop_last=True)\n",
    "test_transform = compute_test_transform()\n",
    "memory_data = CIFAR10Pair(root='data', train=True, transform=test_transform, download=True)\n",
    "memory_loader = DataLoader(memory_data, batch_size=batch_size, shuffle=False, num_workers=16, pin_memory=True)\n",
    "test_data = CIFAR10Pair(root='data', train=False, transform=test_transform, download=True)\n",
    "test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False, num_workers=16, pin_memory=True)\n",
    "\n",
    "# Set up the model and optimizer config.\n",
    "model = Model(feature_dim)\n",
    "model.load_state_dict(torch.load(pretrained_path, map_location='cpu'), strict=False)\n",
    "model = model.to(device)\n",
    "flops, params = profile(model, inputs=(torch.randn(1, 3, 32, 32).to(device),))\n",
    "flops, params = clever_format([flops, params])\n",
    "print('# Model Params: {} FLOPs: {}'.format(params, flops))\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-6)\n",
    "c = len(memory_data.classes)\n",
    "\n",
    "# Training loop.\n",
    "results = {'train_loss': [], 'test_acc@1': [], 'test_acc@5': []} #<< -- output\n",
    "\n",
    "if not os.path.exists('results'):\n",
    "    os.mkdir('results')\n",
    "best_acc = 0.0\n",
    "for epoch in range(1, epochs + 1):\n",
    "    train_loss = train(model, train_loader, optimizer, epoch, epochs, batch_size=batch_size, temperature=temperature, device=device)\n",
    "    results['train_loss'].append(train_loss)\n",
    "    test_acc_1, test_acc_5 = test(model, memory_loader, test_loader, epoch, epochs, c, k=k, temperature=temperature, device=device)\n",
    "    results['test_acc@1'].append(test_acc_1)\n",
    "    results['test_acc@5'].append(test_acc_5)\n",
    "    \n",
    "    # Save statistics.\n",
    "    if test_acc_1 > best_acc:\n",
    "        best_acc = test_acc_1\n",
    "        torch.save(model.state_dict(), './pretrained_model/trained_simclr_model.pth')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Bm8P70VD3mST"
   },
   "source": [
    "# Finetune a Linear Layer for Classification!\n",
    "\n",
    "Now it's time to put the representation vectors to the test!\n",
    "\n",
    "We remove the projection head from the SimCLR model and slap on a linear layer to finetune for a simple classification task. All layers before the linear layer are frozen, and only the weights in the final linear layer are trained. We compare the performance of the SimCLR + finetuning model against a baseline model, where no self-supervised learning is done beforehand, and all weights in the model are trained. You will get to see for yourself the power of self-supervised learning and how the learned representation vectors improve downstream task performance.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Baseline: Without Self-Supervised Learning\n",
    "First, let's take a look at the baseline model. We'll remove the projection head from the SimCLR model and slap on a linear layer to finetune for a simple classification task. No self-supervised learning is done beforehand, and all weights in the model are trained. Run the following cells.\n",
    " \n",
    "**NOTE:** Don't worry if you see low but reasonable performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "is_executing": false
    }
   },
   "outputs": [],
   "source": [
    "class Classifier(nn.Module):\n",
    "    def __init__(self, num_class):\n",
    "        super(Classifier, self).__init__()\n",
    "\n",
    "        # Encoder.\n",
    "        self.f = Model().f\n",
    "        \n",
    "        # Classifier.\n",
    "        self.fc = nn.Linear(2048, num_class, bias=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.f(x)\n",
    "        feature = torch.flatten(x, start_dim=1)\n",
    "        out = self.fc(feature)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Do not modify this cell.\n",
    "feature_dim = 128\n",
    "temperature = 0.5\n",
    "k = 200\n",
    "batch_size = 128\n",
    "epochs = 10\n",
    "percentage = 0.1\n",
    "\n",
    "train_transform = compute_train_transform()\n",
    "train_data = CIFAR10(root='data', train=True, transform=train_transform, download=True)\n",
    "trainset = torch.utils.data.Subset(train_data, list(np.arange(int(len(train_data)*percentage))))\n",
    "train_loader = DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=16, pin_memory=True)\n",
    "test_transform = compute_test_transform()\n",
    "test_data = CIFAR10(root='data', train=False, transform=test_transform, download=True)\n",
    "test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False, num_workers=16, pin_memory=True)\n",
    "\n",
    "model = Classifier(num_class=len(train_data.classes)).to(device)\n",
    "for param in model.f.parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "flops, params = profile(model, inputs=(torch.randn(1, 3, 32, 32).to(device),))\n",
    "flops, params = clever_format([flops, params])\n",
    "print('# Model Params: {} FLOPs: {}'.format(params, flops))\n",
    "optimizer = optim.Adam(model.fc.parameters(), lr=1e-3, weight_decay=1e-6)\n",
    "no_pretrain_results = {'train_loss': [], 'train_acc@1': [], 'train_acc@5': [],\n",
    "           'test_loss': [], 'test_acc@1': [], 'test_acc@5': []}\n",
    "\n",
    "best_acc = 0.0\n",
    "for epoch in range(1, epochs + 1):\n",
    "    train_loss, train_acc_1, train_acc_5 = train_val(model, train_loader, optimizer, epoch, epochs, device='cuda')\n",
    "    no_pretrain_results['train_loss'].append(train_loss)\n",
    "    no_pretrain_results['train_acc@1'].append(train_acc_1)\n",
    "    no_pretrain_results['train_acc@5'].append(train_acc_5)\n",
    "    test_loss, test_acc_1, test_acc_5 = train_val(model, test_loader, None, epoch, epochs)\n",
    "    no_pretrain_results['test_loss'].append(test_loss)\n",
    "    no_pretrain_results['test_acc@1'].append(test_acc_1)\n",
    "    no_pretrain_results['test_acc@5'].append(test_acc_5)\n",
    "    if test_acc_1 > best_acc:\n",
    "        best_acc = test_acc_1\n",
    "        \n",
    "# Print the best test accuracy.\n",
    "print('Best top-1 accuracy without self-supervised learning: ', best_acc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## With Self-Supervised Learning\n",
    "\n",
    "Let's see how much improvement we get with self-supervised learning. Here, we pretrain the SimCLR model using the simclr loss you wrote, remove the projection head from the SimCLR model, and use a linear layer to finetune for a simple classification task."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ssl_best_accuracy"
   },
   "outputs": [],
   "source": [
    "# Do not modify this cell.\n",
    "feature_dim = 128\n",
    "temperature = 0.5\n",
    "k = 200\n",
    "batch_size = 128\n",
    "epochs = 10\n",
    "percentage = 0.1\n",
    "pretrained_path = './pretrained_model/trained_simclr_model.pth'\n",
    "\n",
    "train_transform = compute_train_transform()\n",
    "train_data = CIFAR10(root='data', train=True, transform=train_transform, download=True)\n",
    "trainset = torch.utils.data.Subset(train_data, list(np.arange(int(len(train_data)*percentage))))\n",
    "train_loader = DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=16, pin_memory=True)\n",
    "test_transform = compute_test_transform()\n",
    "test_data = CIFAR10(root='data', train=False, transform=test_transform, download=True)\n",
    "test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False, num_workers=16, pin_memory=True)\n",
    "\n",
    "model = Classifier(num_class=len(train_data.classes))\n",
    "model.load_state_dict(torch.load(pretrained_path, map_location='cpu'), strict=False)\n",
    "model = model.to(device)\n",
    "for param in model.f.parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "flops, params = profile(model, inputs=(torch.randn(1, 3, 32, 32).to(device),))\n",
    "flops, params = clever_format([flops, params])\n",
    "print('# Model Params: {} FLOPs: {}'.format(params, flops))\n",
    "optimizer = optim.Adam(model.fc.parameters(), lr=1e-3, weight_decay=1e-6)\n",
    "pretrain_results = {'train_loss': [], 'train_acc@1': [], 'train_acc@5': [],\n",
    "           'test_loss': [], 'test_acc@1': [], 'test_acc@5': []}\n",
    "\n",
    "best_acc = 0.0\n",
    "for epoch in range(1, epochs + 1):\n",
    "    train_loss, train_acc_1, train_acc_5 = train_val(model, train_loader, optimizer, epoch, epochs)\n",
    "    pretrain_results['train_loss'].append(train_loss)\n",
    "    pretrain_results['train_acc@1'].append(train_acc_1)\n",
    "    pretrain_results['train_acc@5'].append(train_acc_5)\n",
    "    test_loss, test_acc_1, test_acc_5 = train_val(model, test_loader, None, epoch, epochs)\n",
    "    pretrain_results['test_loss'].append(test_loss)\n",
    "    pretrain_results['test_acc@1'].append(test_acc_1)\n",
    "    pretrain_results['test_acc@5'].append(test_acc_5)\n",
    "    if test_acc_1 > best_acc:\n",
    "        best_acc = test_acc_1\n",
    "    \n",
    "# Print the best test accuracy. You should see a best top-1 accuracy of >=70%.\n",
    "print('Best top-1 accuracy with self-supervised learning: ', best_acc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot your Comparison\n",
    "\n",
    "Plot the test accuracies between the baseline model (no pretraining) and same model pretrained with self-supervised learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(no_pretrain_results['test_acc@1'], label=\"Without Pretrain\")\n",
    "plt.plot(pretrain_results['test_acc@1'], label=\"With Pretrain\")\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.title('Test Top-1 Accuracy')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 4
}
